import asyncio

import nest_asyncio
from dotenv import load_dotenv
from langchain.chat_models import init_chat_model
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda

from src.Basic.tools.generatePDF import generate_pdf
from src.Basic.tools.summaryWebSize import summary_website


def main():
    load_dotenv(override=True)

    #定义模型
    model = init_chat_model(model="deepseek-chat", model_provider="deepseek")
    # 定义 天气查询 工具函数
    optimized_prompt = ChatPromptTemplate.from_template(
        """
        请优化以下网站总结的内容，使其更适合PDF报告格式：
        网站总结：
        {summary}
        
        请重新组织内容，包括：
        1.清晰的标题和结构
        2.要点总结
        3.详细说明
        4.使用林妹妹等
        
        优化后的内容：
        """
    )

    def debug_print(x):
        print("中间结果（新闻正文）:", x)
        return x

    def arrayToStr(x):
        return x[0]

    debug_node = RunnableLambda(debug_print)
    exchange_node = RunnableLambda(arrayToStr)

    while True:
        simple_chain = summary_website | debug_node | exchange_node | generate_pdf
        optimized_chain = (summary_website
                           | (lambda summary: {"summary": summary})
                           | optimized_prompt
                           | model
                           | StrOutputParser()
                           | generate_pdf
                           )

        url = input("请输入要查询的网址:")
        if url == "exit":
            break

        use_optimized_chain = input("请选择使用优化的链（1）还是简单的链（2）:")

        if use_optimized_chain == "2":
            chain = simple_chain
        else:
            chain = optimized_chain

        chain.invoke({"url": url})  #


if __name__ == "__main__":
    main()